import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine

# 数据库配置
db_conf = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
}

# 创建数据库引擎
engine = create_engine(
    f"mysql+pymysql://{db_conf['user']}:{db_conf['password']}@{db_conf['host']}:{db_conf['port']}/{db_conf['database']}"
)

# SQL 查询
query = """
SELECT d.* FROM date_1 d WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
"""

# 使用 SQLAlchemy 引擎执行查询并分块读取数据
chunk_size = 10000
df_chunks = pd.read_sql_query(query, engine, chunksize=chunk_size)
df1 = pd.concat(df_chunks, ignore_index=True)

# 计算 zd_closes 列
df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)

# 打印前几行数据
print(df1.head())

# 删除包含 NaN 值的行
df1 = df1.dropna(subset=['zd_closes'])

# 再次打印前几行数据
print(df1.head())

# 定义要排除的列
ex = ['id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low', 'closes', 'pre_closes', 'changes', 'pct_chg', 'amount']

# 获取数值类型的列，并去除指定的列
number = df1.select_dtypes(include=['number']).columns.tolist()
newList = [col for col in number if col not in ex]

# 构建回归公式
formuls = 'zd_closes ~ ' + ' + '.join(newList)

# 拟合 OLS 回归模型
res = smf.ols(formuls, data=df1).fit()

# 输出回归结果摘要
print(res.summary())